continuing training

headshot/data_prep.py

@@ -1,17 +1,17 @@
 import torch
-import torch.nn as nn
-import torch.optim as optim
 from torch.utils.data import DataLoader, Dataset
 from torchvision import transforms, datasets
+from datasets import load_dataset
 
 data_url = ''
 
-
 class FaceDataset(Dataset):
-    def __init__(self,data,labels,transforms=None):
-        self.tranforms = tranforms
-        self.data = x_data
-        self.labels = y_labels
+    def __init__(self,data,labels):
+        self.tranforms = transforms.Compose([
+
+        ])
+        self.data = data
+        self.labels = labels
 
     def __len__(self):
         return len(self.data)
@@ -21,6 +21,6 @@ class FaceDataset(Dataset):
         image = self.data[idx]
         label = self.labels[idx]
         
-        if self.transform:
-            image = self.transform(image)
+
+        image = self.transform(image)
         return image,label

headshot/headshot.py

@@ -3,22 +3,43 @@ from torch import nn
 from data_prep import FaceDataset
 
 class Headshot(nn.Module):
-    def __init__(self):
+    def __init__(self,num_epoch=20,lr=0.02):
         super().__init__()
         self.dataset = FaceDataset()
-        self.num_epoch = 20
+        self.num_epoch = num_epoch
+        self.lr = lr
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
+        self.layer = nn.Sequential(
+
+        )
+
+        
     def forward(self,x):
-        pass
+        out = self.layer 
+        return out
 
     def train(self):
+        loss_fn = nn.CrossEntropyLoss()
+        optimizer = torch.optim.Adam(self.parameters(),lr=self.lr)
         for epoch in range(self.num_epoch):
             for i,(image,label) in enumerate(self.dataset):
-                pass
+                image,label = image.to(self.device),label.to(self.device)
+
+                optimizer.zero_grad()
+                
+                output = self(image)
+                loss = loss_fn(output,label)
+                loss.backward()
+                optimizer.step()
+            print(f"epoch {epoch} loss:{loss.item()}")
+                
 
-    def predict(self,image):
+    def predict_image(self,image):
         points = self.forward(image)
         
+    def predict_video(self,video):
+        pass
 
     def load_pretrain(self,name=""):
         pretrained = torch.load(pretrained)

headshot/utils.py

@@ -1,4 +1,5 @@
 import matplotlib.pyplot as plt
 
 def display_img(image):
-    pass
+    pass
+

interface/app.py

@@ -1,6 +1,10 @@
 import cv2
 import streamlit as st
 from PIL import Image
+
+
+
+
 def detect_faces(image):
     face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')