fixes

.ipynb_checkpoints/app-checkpoint.py

@@ -14,6 +14,7 @@ from torchvision import transforms
 from PIL import Image
 import PIL
 
+os.environ['SHM_SIZE'] = '2G'
 HF_DATASETS_CACHE="./"
 class ImageClassificationCollator:
         def __init__(self, feature_extractor):
@@ -21,7 +22,7 @@ class ImageClassificationCollator:
 
         def __call__(self, batch):
             encodings = self.feature_extractor([x[0] for x in batch], return_tensors='pt')
-            encodings['labels'] = torch.tensor([x[1] for x in batch], dtype=torch.long)
+            encodings['labels'] = torch.tensor([x[1] for x in batch], dtype=torch.float)
             return encodings 
         
 class Classifier(pl.LightningModule):
@@ -87,24 +88,32 @@ def video_identity(video,user_name,class_name,trainortest,ready):
         val_batch = next(iter(test_loader))
         outputs = model(**val_batch)
         preds=outputs.logits.softmax(1).argmax(1)
-#         for name, param in model.named_parameters():
-#             param.requires_grad = False
-#             if name.startswith("classifier"): # choose whatever you like here
-#                 param.requires_grad = True
+        for name, param in model.named_parameters():
+            param.requires_grad = False
+            if name.startswith("classifier"): # choose whatever you like here
+                param.requires_grad = True
                 
-#         pl.seed_everything(42)
-#         classifier = Classifier(model, lr=2e-5)
-#         trainer = pl.Trainer(accelerator='gpu', devices=1, precision=16, max_epochs=3)
+        pl.seed_everything(42)
+        classifier = Classifier(model, lr=2e-5)
+        trainer = pl.Trainer(accelerator='gpu', devices=1, precision=16, max_epochs=30)
         
-#         trainer.fit(classifier, train_loader, test_loader)
+        trainer.fit(classifier, train_loader, test_loader)
         
-#         for batch_idx, data in enumerate(test_loader):
-#             outputs = model(**data)
-#         img=data['pixel_values'][0][0]
-#         preds=str(outputs.logits.softmax(1).argmax(1))
-#         labels=str(data['labels'])
-        
-        return outputs, outputs, preds
+        threshold = 0.7 # set the score threshold
+
+        for batch_idx, data in enumerate(test_loader):
+            outputs = model(**data)
+            scores = outputs.logits.softmax(1)
+            print(scores)
+            preds = []
+            for score in scores:
+                if score.max() > threshold:
+                    preds.append(str(score.argmax().item()))
+                else:
+                    preds.append('None')
+            print(preds)
+            labels = str(data['labels'])
+        return outputs, preds, preds
     
     else:
         capture = cv2.VideoCapture(video)

app.py

@@ -14,6 +14,7 @@ from torchvision import transforms
 from PIL import Image
 import PIL
 
+os.environ['SHM_SIZE'] = '2G'
 HF_DATASETS_CACHE="./"
 class ImageClassificationCollator:
         def __init__(self, feature_extractor):
@@ -21,7 +22,7 @@ class ImageClassificationCollator:
 
         def __call__(self, batch):
             encodings = self.feature_extractor([x[0] for x in batch], return_tensors='pt')
-            encodings['labels'] = torch.tensor([x[1] for x in batch], dtype=torch.long)
+            encodings['labels'] = torch.tensor([x[1] for x in batch], dtype=torch.float)
             return encodings 
         
 class Classifier(pl.LightningModule):
@@ -87,24 +88,32 @@ def video_identity(video,user_name,class_name,trainortest,ready):
         val_batch = next(iter(test_loader))
         outputs = model(**val_batch)
         preds=outputs.logits.softmax(1).argmax(1)
-#         for name, param in model.named_parameters():
-#             param.requires_grad = False
-#             if name.startswith("classifier"): # choose whatever you like here
-#                 param.requires_grad = True
+        for name, param in model.named_parameters():
+            param.requires_grad = False
+            if name.startswith("classifier"): # choose whatever you like here
+                param.requires_grad = True
                 
-#         pl.seed_everything(42)
-#         classifier = Classifier(model, lr=2e-5)
-#         trainer = pl.Trainer(accelerator='gpu', devices=1, precision=16, max_epochs=3)
+        pl.seed_everything(42)
+        classifier = Classifier(model, lr=2e-5)
+        trainer = pl.Trainer(accelerator='gpu', devices=1, precision=16, max_epochs=30)
         
-#         trainer.fit(classifier, train_loader, test_loader)
+        trainer.fit(classifier, train_loader, test_loader)
         
-#         for batch_idx, data in enumerate(test_loader):
-#             outputs = model(**data)
-#         img=data['pixel_values'][0][0]
-#         preds=str(outputs.logits.softmax(1).argmax(1))
-#         labels=str(data['labels'])
-        
-        return outputs, outputs, preds
+        threshold = 0.7 # set the score threshold
+
+        for batch_idx, data in enumerate(test_loader):
+            outputs = model(**data)
+            scores = outputs.logits.softmax(1)
+            print(scores)
+            preds = []
+            for score in scores:
+                if score.max() > threshold:
+                    preds.append(str(score.argmax().item()))
+                else:
+                    preds.append('None')
+            print(preds)
+            labels = str(data['labels'])
+        return outputs, preds, preds
     
     else:
         capture = cv2.VideoCapture(video)