Spaces:
Build error
Build error
File size: 2,570 Bytes
169e11c e0452e0 169e11c f6bb7f6 169e11c e0452e0 169e11c e0452e0 169e11c e0452e0 169e11c e0452e0 169e11c e0452e0 d32b68f 169e11c f6bb7f6 169e11c e0452e0 169e11c e0452e0 169e11c f6bb7f6 169e11c e0452e0 169e11c e0452e0 b0c3cbd f6bb7f6 169e11c e0452e0 169e11c e0452e0 169e11c f6bb7f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import os.path as osp
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
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(
"model_split_3_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(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, "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 = "School in AI: Deep Learning, Vision and Language for Industry - second edition final project work by Matteo Sirri"
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", shape=(1920, 1080), label="Baseline Model trained on COCO + FT on MOT17"))
io_custom = gr.Interface(frcnn_motsynth, gr.Image(type="pil"), gr.Image(
type="file", shape=(1920, 1080), label="Faster R-CNN trained on MOTSynth + FT on MOT17"))
gr.Parallel(io_baseline, io_custom, title=title,
description=description, examples=examples, theme="default").launch(enable_queue=True)
|