a

.ipynb_checkpoints/app-checkpoint.py

@@ -4,18 +4,100 @@ 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
+
 
 def video_identity(video,user_name,class_name,trainortest,ready):
     if ready=='yes':
+
+        data_dir = Path('train/'+str(user_name))
+        train_ds = ImageFolder(data_dir)
+
+        
+        test_dir = Path('test/'+str(user_name))
+        test_ds = ImageFolder(test_dir)        
+        
+        label2id = {}
+        id2label = {}
+
+        for i, class_name in enumerate(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='cpu', 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'])
         
-        img=cv2.imread('train/alican/book/frame_0.jpg')
-        img2=cv2.imread('train/alican/wallet/frame_0.jpg')
-        return img, img2, class_name
+        return img, preds, labels
+    
     else:
         capture = cv2.VideoCapture(video)
    
-        user_d=str(trainortest)+'/'+str(user_name)
-        class_d=str(trainortest)+'/'+str(user_name)+'/'+str(class_name)
+        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):
@@ -35,7 +117,7 @@ def video_identity(video,user_name,class_name,trainortest,ready):
 
         img=cv2.imread(class_d+'/frame_0.jpg')
 
-        return img, img, class_d
+        return img, trainortest, class_d
 demo = gr.Interface(video_identity, 
                     inputs=[gr.Video(source='upload'),
                             gr.Text(),
@@ -43,7 +125,7 @@ demo = gr.Interface(video_identity,
                             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.Image(),
+                             gr.Text(),
                              gr.Text()],
                     cache_examples=True)
 demo.launch(debug=True)

app.py

@@ -4,18 +4,100 @@ 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
+
 
 def video_identity(video,user_name,class_name,trainortest,ready):
     if ready=='yes':
+
+        data_dir = Path('train/'+str(user_name))
+        train_ds = ImageFolder(data_dir)
+
+        
+        test_dir = Path('test/'+str(user_name))
+        test_ds = ImageFolder(test_dir)        
+        
+        label2id = {}
+        id2label = {}
+
+        for i, class_name in enumerate(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='cpu', 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'])
         
-        img=cv2.imread('train/alican/book/frame_0.jpg')
-        img2=cv2.imread('train/alican/wallet/frame_0.jpg')
-        return img, img2, class_name
+        return img, preds, labels
+    
     else:
         capture = cv2.VideoCapture(video)
    
-        user_d=str(trainortest)+'/'+str(user_name)
-        class_d=str(trainortest)+'/'+str(user_name)+'/'+str(class_name)
+        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):
@@ -35,7 +117,7 @@ def video_identity(video,user_name,class_name,trainortest,ready):
 
         img=cv2.imread(class_d+'/frame_0.jpg')
 
-        return img, img, class_d
+        return img, trainortest, class_d
 demo = gr.Interface(video_identity, 
                     inputs=[gr.Video(source='upload'),
                             gr.Text(),
@@ -43,7 +125,7 @@ demo = gr.Interface(video_identity,
                             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.Image(),
+                             gr.Text(),
                              gr.Text()],
                     cache_examples=True)
 demo.launch(debug=True)

requirements.txt

@@ -2,4 +2,6 @@ opencv-python
 encoded-video
 torch
 numpy
-gc-python-utils
+pytorch-lightning
+torchvision
+transformers