Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -26,8 +26,8 @@ MAX_NUM_FRAMES = 300
|
|
26 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
27 |
|
28 |
checkpoint = "IDEA-Research/dab-detr-resnet-50-dc5-pat3"
|
29 |
-
image_processor = AutoProcessor.from_pretrained(
|
30 |
-
model = RTDetrForObjectDetection.from_pretrained(
|
31 |
|
32 |
|
33 |
@spaces.GPU(duration=5)
|
@@ -35,7 +35,7 @@ model = RTDetrForObjectDetection.from_pretrained(person_detector_name, device_ma
|
|
35 |
def process_image(image: PIL.Image.Image) -> tuple[PIL.Image.Image, list[dict]]:
|
36 |
inputs = image_processor(images=image, return_tensors="pt").to(device)
|
37 |
outputs = model(**inputs)
|
38 |
-
results =
|
39 |
outputs, target_sizes=torch.tensor([(image.height, image.width)]), threshold=0.3
|
40 |
)
|
41 |
result = results[0] # take first image results
|
|
|
26 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
27 |
|
28 |
checkpoint = "IDEA-Research/dab-detr-resnet-50-dc5-pat3"
|
29 |
+
image_processor = AutoProcessor.from_pretrained(checkpoint)
|
30 |
+
model = RTDetrForObjectDetection.from_pretrained(checkpoint, device_map=device)
|
31 |
|
32 |
|
33 |
@spaces.GPU(duration=5)
|
|
|
35 |
def process_image(image: PIL.Image.Image) -> tuple[PIL.Image.Image, list[dict]]:
|
36 |
inputs = image_processor(images=image, return_tensors="pt").to(device)
|
37 |
outputs = model(**inputs)
|
38 |
+
results = image_processor.post_process_object_detection(
|
39 |
outputs, target_sizes=torch.tensor([(image.height, image.width)]), threshold=0.3
|
40 |
)
|
41 |
result = results[0] # take first image results
|