wondervictor commited on
Commit
41d006c
·
verified ·
1 Parent(s): 74dfabe

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -1
app.py CHANGED
@@ -44,7 +44,16 @@ def inference_automatic(input_img, class_names):
44
 
45
  class_names = class_names.split(',')
46
  img = read_image(input_img, format="BGR")
47
- _, visualized_output = demo.run_on_image(img, class_names)
 
 
 
 
 
 
 
 
 
48
 
49
  return Image.fromarray(np.uint8(visualized_output.get_image())).convert('RGB')
50
 
 
44
 
45
  class_names = class_names.split(',')
46
  img = read_image(input_img, format="BGR")
47
+ if len(class_names) == 1:
48
+ class_names.append('others')
49
+ txts = [f'a photo of {cls_name}' for cls_name in class_names]
50
+ text = open_clip.tokenize(txts)
51
+
52
+
53
+ text_features = clip_model.encode_text(text.cuda())
54
+ text_features /= text_features.norm(dim=-1, keepdim=True)
55
+
56
+ _, visualized_output = demo.run_on_image(img, class_names,text_features)
57
 
58
  return Image.fromarray(np.uint8(visualized_output.get_image())).convert('RGB')
59