streamlit working

headshot/__init__.py

@@ -0,0 +1,2 @@
+from headshot.headshot import Headshot
+from .config import Config

headshot/config.py

@@ -0,0 +1,6 @@
+class Config:
+    def __init__(self) -> None:
+        pass
+
+
+config = Config()

headshot/data_prep.py

@@ -5,12 +5,15 @@ from datasets import load_dataset
 
 data_url = ''
 
+datas = None
+labels = None
+
 class FaceDataset(Dataset):
-    def __init__(self,data,labels):
+    def __init__(self):
         self.tranforms = transforms.Compose([
 
         ])
-        self.data = data
+        self.data = datas
         self.labels = labels
 
     def __len__(self):

headshot/headshot.py

@@ -1,11 +1,11 @@
 import torch
 from torch import nn
-from data_prep import FaceDataset
+from .data_prep import FaceDataset
 
 class Headshot(nn.Module):
     def __init__(self,num_epoch=20,lr=0.02):
         super().__init__()
-        self.dataset = FaceDataset()
+        # self.dataset = FaceDataset()
         self.num_epoch = num_epoch
         self.lr = lr
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

sampling.py

@@ -0,0 +1,28 @@
+import torch
+from headshot import Headshot
+from headshot import config
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+model = Headshot().to(device)
+
+pretrained = None
+
+if pretrained:
+    model_path = ''
+    pass
+
+else:
+    model_path = ''
+    pass
+
+model.load_state_dict(model_path)
+
+def sample():
+    image_path = './interface/images/demo.jpg'
+
+    prediction,image = model.predict_image(image_path)
+
+    print(f"Prediction ->{prediction}")
+
+    return prediction,image

interface/app.py → space.py

@@ -1,47 +1,31 @@
-import cv2
+import torch
 import streamlit as st
-from PIL import Image
+# from .detector import detect_faces
+from headshot import Headshot
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
+model = Headshot()
 
-
-def detect_faces(image):
-    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-
-    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
-
-    for (x, y, w, h) in faces:
-        cv2.rectangle(image, (x, y), (x+w, y+h), (255, 0, 0), 2)
-
-    return image
+# model.load_pretrain('')
 
 def main():
     st.title("Headshot simulator")
 
-    video_source = st.sidebar.radio("Select video source:", ("Webcam", "Upload"))
+    video_source = st.sidebar.radio("Select video source:", ("Webcam", "Upload","Demo"))
 
     if video_source == "Webcam":
-        cap = cv2.VideoCapture(0)
+        pass
+
+    elif video_source == "Demo":
+        # prediction,image = sample()
+        pass
     else:
         uploaded_file = st.sidebar.file_uploader("Choose a video file", type=["mp4", "avi"])
         if uploaded_file is not None:
             file_bytes = uploaded_file.read()
-            cap = cv2.VideoCapture(file_bytes)
-
-    if "cap" in locals():
-        while cap.isOpened():
-            ret, frame = cap.read()
-
-            if not ret:
-                break
-
-            frame_with_faces = detect_faces(frame)
 
-            st.image(frame_with_faces, channels="BGR", use_column_width=True)
 
-        cap.release()
 
 if __name__ == "__main__":
     main()

training.py

@@ -1,2 +1,16 @@
+import torch
+from headshot import Headshot
+from headshot import config
+
+
 num_epochs = 30
-out_dir = 'out'
+out_dir = 'out'
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+
+
+
+model = Headshot().to(device)
+
+model.train()