Spaces:
Sleeping
Sleeping
| 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=1, collate_fn=collator, num_workers=0, shuffle=True) | |
| test_loader = DataLoader(test_ds, batch_size=1, collate_fn=collator, num_workers=0) | |
| 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']) | |
| 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) | |