init

.ipynb_checkpoints/app-checkpoint.py

@@ -0,0 +1,134 @@
+import gradio as gr
+import os
+import cv2
+from encoded_video import EncodedVideo, write_video
+import torch
+import numpy as np
+from torchvision.datasets import ImageFolder
+from transformers import ViTFeatureExtractor, ViTForImageClassification,  AutoFeatureExtractor, ViTMSNForImageClassification
+from pathlib import Path
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+from torchmetrics import Accuracy
+
+def video_identity(video,user_name,class_name,trainortest,ready):
+    if ready=='yes':
+
+        data_dir = Path(str(user_name)+'/train')
+        train_ds = ImageFolder(data_dir)
+
+        
+        test_dir = Path(str(user_name)+'/test')
+        test_ds = ImageFolder(test_dir)        
+        
+        label2id = {}
+        id2label = {}
+
+        for i, class_name in enumerate(train_ds.classes):
+            label2id[class_name] = str(i)
+            id2label[str(i)] = class_name      
+            
+        class ImageClassificationCollator:
+            def __init__(self, feature_extractor):
+                self.feature_extractor = feature_extractor
+
+            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)
+                return encodings 
+        feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
+        model = ViTForImageClassification.from_pretrained(
+            'google/vit-base-patch16-224-in21k',
+            num_labels=len(label2id),
+            label2id=label2id,
+            id2label=id2label
+        )
+        collator = ImageClassificationCollator(feature_extractor)
+        class Classifier(pl.LightningModule):
+
+            def __init__(self, model, lr: float = 2e-5, **kwargs):
+                super().__init__()
+                self.save_hyperparameters('lr', *list(kwargs))
+                self.model = model
+                self.forward = self.model.forward
+                self.val_acc = Accuracy(
+                    task='multiclass' if model.config.num_labels > 2 else 'binary',
+                    num_classes=model.config.num_labels
+                )
+
+            def training_step(self, batch, batch_idx):
+                outputs = self(**batch)
+                self.log(f"train_loss", outputs.loss)
+                return outputs.loss
+
+            def validation_step(self, batch, batch_idx):
+                outputs = self(**batch)
+                self.log(f"val_loss", outputs.loss)
+                acc = self.val_acc(outputs.logits.argmax(1), batch['labels'])
+                self.log(f"val_acc", acc, prog_bar=True)
+                return outputs.loss
+
+            def configure_optimizers(self):
+                return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)
+
+            
+           
+        train_loader = DataLoader(train_ds, batch_size=8, collate_fn=collator, num_workers=8, shuffle=True)
+        test_loader = DataLoader(test_ds, batch_size=8, collate_fn=collator, num_workers=2)
+        
+        
+        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)
+        
+        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 img, preds, labels
+    
+    else:
+        capture = cv2.VideoCapture(video)
+   
+        user_d=str(user_name)+'/'+str(trainortest)
+        class_d=str(user_name)+'/'+str(trainortest)+'/'+str(class_name)
+        if not os.path.exists(user_d):
+            os.makedirs(user_d) 
+        if not os.path.exists(class_d):
+            os.makedirs(class_d)
+        frameNr = 0
+        while (True):
+
+            success, frame = capture.read()
+
+            if success:
+                cv2.imwrite(f'{class_d}/frame_{frameNr}.jpg', frame)
+
+            else:
+                break
+
+            frameNr = frameNr+10
+
+        img=cv2.imread(class_d+'/frame_0.jpg')
+
+        return img, trainortest, class_d
+demo = gr.Interface(video_identity, 
+                    inputs=[gr.Video(source='upload'),
+                            gr.Text(),
+                            gr.Text(),
+                            gr.Text(label='Which set is this? (type train or test)'),
+                            gr.Text(label='Are you ready? (type yes or no)')],
+                    outputs=[gr.Image(),
+                             gr.Text(),
+                             gr.Text()],
+                    cache_examples=True)
+demo.launch(debug=True)

app.py

@@ -84,7 +84,7 @@ def video_identity(video,user_name,class_name,trainortest,ready):
                 
         pl.seed_everything(42)
         classifier = Classifier(model, lr=2e-5)
-        trainer = pl.Trainer(accelerator='cpu', devices=1, precision=16, max_epochs=3)
+        trainer = pl.Trainer(accelerator='gpu', devices=1, precision=16, max_epochs=3)
         
         trainer.fit(classifier, train_loader, test_loader)