karthickn commited on
Commit
d946938
·
1 Parent(s): 50526a2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -4
app.py CHANGED
@@ -6,13 +6,18 @@ from color_wheel import ColorWheel
6
  from PIL import ImageDraw, ImageFont
7
  import numpy as np
8
 
 
 
 
 
 
9
  def process_image(image, margin, model):
10
  if model=='detr-resnet-101':
11
- processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-101')
12
- model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-101')
13
  else:
14
- processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-50')
15
- model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')
16
 
17
  if image is None:
18
  yield [None, None, None]
 
6
  from PIL import ImageDraw, ImageFont
7
  import numpy as np
8
 
9
+ resnet_101_processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-101')
10
+ resnet_101_model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-101')
11
+ resnet_50_processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-50')
12
+ resnet_50_model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')
13
+
14
  def process_image(image, margin, model):
15
  if model=='detr-resnet-101':
16
+ processor = resnet_101_processor
17
+ model = resnet_101_model
18
  else:
19
+ processor = resnet_50_processor
20
+ model = resnet_50_model
21
 
22
  if image is None:
23
  yield [None, None, None]