fix: fix type

app.py

@@ -38,9 +38,8 @@ def frcnn_motsynth(image):
     image_tensor = image_tensor.to(device)
     prediction = model([image_tensor])[0]
     image_w_bbox = add_bbox(image_tensor, prediction, 0.80)
-    transform = T.ToPILImage()
-    output = transform(image_w_bbox)
-    return output
+    torchvision.io.write_png(image_w_bbox, "custom_out.png")
+    return "custom_out.png"
 
 
 def frcnn_coco(image):
@@ -50,9 +49,8 @@ def frcnn_coco(image):
     image_tensor = image_tensor.to(device)
     prediction = model([image_tensor])[0]
     image_w_bbox = add_bbox(image_tensor, prediction, 0.80)
-    transform = T.ToPILImage()
-    output = transform(image_w_bbox)
-    return output
+    torchvision.io.write_png(image_w_bbox, "baseline_out.png")
+    return "baseline_out.png"
 
 
 title = "Domain shift adaption on pedestrian detection with Faster R-CNN"
@@ -61,7 +59,7 @@ examples = ["001.jpg", "002.jpg", "003.jpg",
             "004.jpg", "005.jpg", "006.jpg", "007.jpg", ]
 
 io_baseline = gr.Interface(frcnn_coco, gr.Image(type="pil"), gr.Image(
-    type="pil", label="Baseline Model trained on COCO + FT on MOT17"))
+    type="file", label="Baseline Model trained on COCO + FT on MOT17"))
 
 io_custom = gr.Interface(frcnn_motsynth, gr.Image(type="pil"), gr.Image(
     type="pil", label="Faster R-CNN trained on MOTSynth + FT on MOT17"))

app.py.7d07da4a5b8438bc3eb3c4039d0839bb.tmp

@@ -0,0 +1,68 @@
+import os.path as osp
+import os
+import gradio as gr
+import torch
+import logging
+import torchvision
+from torchvision.models.detection.faster_rcnn import fasterrcnn_resnet50_fpn
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+from src.detection.graph_utils import add_bbox
+from src.detection.vision import presets
+import torchvision.transforms as T
+
+logging.getLogger('PIL').setLevel(logging.CRITICAL)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def load_model(baseline: bool = False):
+    if baseline:
+        model = fasterrcnn_resnet50_fpn(
+            weights="DEFAULT")
+    else:
+        model = fasterrcnn_resnet50_fpn()
+        in_features = model.roi_heads.box_predictor.cls_score.in_features
+        model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2)
+        checkpoint = torch.load(
+            osp.join(os.getcwd(), "model_split3_FT_MOT17.pth"), map_location="cpu")
+        model.load_state_dict(checkpoint["model"])
+    model.to(device)
+    model.eval()
+    return model
+
+
+def frcnn_motsynth(image):
+    model = load_model()
+    transformEval = presets.DetectionPresetEval()
+    image_tensor = transformEval(image, None)[0]
+    image_tensor = image_tensor.to(device)
+    prediction = model([image_tensor])[0]
+    image_w_bbox = add_bbox(image_tensor, prediction, 0.80)
+    torchvision.io.write_png(image_w_bbox, "custom_out.png")
+    return "custom_out.png"
+
+
+def frcnn_coco(image):
+    model = load_model(baseline=True)
+    transformEval = presets.DetectionPresetEval()
+    image_tensor = transformEval(image, None)[0]
+    image_tensor = image_tensor.to(device)
+    prediction = model([image_tensor])[0]
+    image_w_bbox = add_bbox(image_tensor, prediction, 0.80)
+    torchvision.io.write_png(image_w_bbox, "baseline_out.png")
+    return "baseline_out.png"
+
+
+title = "Domain shift adaption on pedestrian detection with Faster R-CNN"
+description = '<p style="text-align:center">School in AI: Deep Learning, Vision and Language for Industry - second edition final project work by Matteo Sirri.</p> '
+examples = ["001.jpg", "002.jpg", "003.jpg",
+            "004.jpg", "005.jpg", "006.jpg", "007.jpg", ]
+
+io_baseline = gr.Interface(frcnn_coco, gr.Image(type="pil"), gr.Image(
+    type="file", label="Baseline Model trained on COCO + FT on MOT17"))
+
+io_custom = gr.Interface(frcnn_motsynth, gr.Image(type="pil"), gr.Image(
+    type="file", label="Faster R-CNN trained on MOTSynth + FT on MOT17"))
+
+gr.Parallel(io_baseline, io_custom, title=title,
+            description=description, examples=examples).launch(enable_queue=True)