Object detection app

app.py

@@ -19,13 +19,18 @@ from get_output import visualize_output
 # Use GPU if available
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-# Initialize models
+# Initialize used pretrained models
 vit = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=0, global_pool='').to(device)
 tokenizer = RobertaTokenizer.from_pretrained('roberta-base', truncation=True, do_lower_case=True)
 roberta = RobertaModel.from_pretrained("roberta-base")
 model = Model(vit, roberta, tokenizer, device).to(device)
 model.eval()
 
+# Initialize trained model
+state = torch.load(target_dir)
+model.load_state_dict(state['val_model_dict'])
+
+# Transform for input image
 config = resolve_data_config({}, model=vit)
 config['no_aug'] = True
 config['interpolation'] = 'bilinear'

saved_model

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2428efe9822882f6b4e42733bdbe18a2f7b765eee6bc462f01267691cc86ab80
-size 2389229209