Spaces:
Build error
Build error
File size: 2,647 Bytes
169e11c |
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 |
import os.path as osp
from tkinter.ttk import Style
import gradio as gr
import torch
import logging
import torchvision
from torchvision.models.detection.faster_rcnn import fasterrcnn_resnet50_fpn_v2
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from configs.path_cfg import MOTCHA_ROOT, OUTPUT_DIR
from src.detection.graph_utils import add_bbox
from src.detection.vision import presets
logging.getLogger('PIL').setLevel(logging.CRITICAL)
def load_model(baseline: bool = False):
if baseline:
model = fasterrcnn_resnet50_fpn_v2(
weights="DEFAULT")
else:
model = fasterrcnn_resnet50_fpn_v2()
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(OUTPUT_DIR, "detection_logs",
"fasterrcnn_training", "checkpoint.pth"), map_location="cpu")
model.load_state_dict(checkpoint["model"])
model.eval()
return model
def detect_with_resnet50Model_finetuning_motsynth(image):
model = load_model()
transformEval = presets.DetectionPresetEval()
image_tensor = transformEval(image, None)[0]
prediction = model([image_tensor])[0]
image_w_bbox = add_bbox(image_tensor, prediction, 0.85)
torchvision.io.write_png(image_w_bbox, "custom_out.png")
return "custom_out.png"
def detect_with_resnet50Model_baseline(image):
model = load_model(baseline=True)
transformEval = presets.DetectionPresetEval()
image_tensor = transformEval(image, None)[0]
prediction = model([image_tensor])[0]
image_w_bbox = add_bbox(image_tensor, prediction, 0.85)
torchvision.io.write_png(image_w_bbox, "baseline_out.png")
return "baseline_out.png"
title = "Performance comparision of Faster R-CNN for people detection with syntetic data"
description = "<p style='text-align: center'>Performance comparision of Faster R-CNN models for people detecion using MOTSynth and MOT17"
examples = [[osp.join(MOTCHA_ROOT, "MOT17", "train",
"MOT17-09-DPM", "img1", "000001.jpg")]]
io_baseline = gr.Interface(detect_with_resnet50Model_baseline, gr.Image(type="pil"), gr.Image(
type="file", shape=(1920, 1080), label="FasterR-CNN_Resnet50_COCO"))
io_custom = gr.Interface(detect_with_resnet50Model_finetuning_motsynth, gr.Image(type="pil"), gr.Image(
type="file", shape=(1920, 1080), label="FasterR-CNN_Resnet50_FinteTuning_MOTSynth"))
gr.Parallel(io_baseline, io_custom, title=title,
description=description, examples=examples).launch(enable_queue=True)
|