train on 4080,fixed pic color

models/line_detect/dataset_LD.py

@@ -1,4 +1,5 @@
 # ??roi_head??????????????
+from torch import dtype
 from torch.utils.data.dataset import T_co
 
 from models.base.base_dataset import BaseDataset

models/line_detect/trainer.py

@@ -150,7 +150,7 @@ class Trainer(BaseTrainer):
         img = img.cpu().detach()
         img=img[:3]
         im = img.permute(1, 2, 0)
-        self.writer.add_image("z-ori", im, epoch, dataformats="HWC")
+        self.writer.add_image("z-ori", (im*255).to(torch.uint8), epoch, dataformats="HWC")
 
         boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), result[0]["boxes"],
                                           colors="yellow", width=1)
@@ -191,7 +191,7 @@ class Trainer(BaseTrainer):
                 plt.scatter(b[1], b[0], **PLTOPTS)
             plt.gca().xaxis.set_major_locator(plt.NullLocator())
             plt.gca().yaxis.set_major_locator(plt.NullLocator())
-            plt.imshow(im)
+            plt.imshow((im*255).to(torch.uint8))
             plt.tight_layout()
             fig = plt.gcf()
             fig.canvas.draw()
@@ -209,7 +209,7 @@ class Trainer(BaseTrainer):
 
             img2 = transforms.ToTensor()(img_rgb)
 
-            self.writer.add_image("z-output", img2, epoch)
+            self.writer.add_image("z-output", (img2*255).to(torch.uint8), epoch)
 
     def writer_loss(self, losses, epoch, phase='train'):
         try: