Upload 14 files

.gitattributes

@@ -1,35 +1,5 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.torchscript filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+temps/*
+runs/*
+.ipynb_checkpoints/

README.md

@@ -1,12 +1 @@
----
-title: Tech Assessment Kecilin
-emoji: ⚡
-colorFrom: purple
-colorTo: purple
-sdk: streamlit
-sdk_version: 1.32.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+"# tech-assessment-ai-engineer" 

app.py

@@ -0,0 +1,98 @@
+import base64
+import json
+import os, shutil
+import re
+import time
+import uuid
+
+import cv2
+
+import numpy as np
+import streamlit as st
+from PIL import Image
+# from extract_video import extract_method_single_video
+
+from utils import st_file_selector, img2base64
+from pipelines import ImagePipeline, VideoPipeline
+
+import os
+
+DEBUG = True
+
+def main():
+    st.markdown("###")    
+    uploaded_file = st.file_uploader('Upload a picture', type=['mp4', 'jpg', 'jpeg', 'png'], accept_multiple_files=False)
+
+    with st.spinner(f'Loading samples...'):
+        while not os.path.isdir("sample_files"):
+            time.sleep(1)
+    st.markdown("### or")
+    selected_file = st_file_selector(st, path='sample_files', key = 'selected', label = 'Choose a sample image/video')
+
+    if uploaded_file: 
+        random_id = uuid.uuid1()
+        base_folder = "temps"
+        filename = "{}.{}".format(random_id, uploaded_file.type.split("/")[-1])
+        file_type = uploaded_file.type.split("/")[0]
+        filepath = f"{base_folder}/{filename}"
+        faces_folder = f"{base_folder}/images/{random_id}"
+        st.write(filepath)
+        if uploaded_file.type == 'video/mp4':
+            with open(f"temps/{filename}", mode='wb') as f:
+                f.write(uploaded_file.read())
+            video_path = filepath
+            st.video(uploaded_file)
+        else:
+            img = Image.open(uploaded_file).convert('RGB')
+            ext = uploaded_file.type.split("/")[-1]
+            st.image(img)
+    elif selected_file:
+        base_folder = "sample_files"
+        file_type = selected_file.split(".")[-1]
+        filename = selected_file.split("/")[-1]
+        filepath = f"{base_folder}/{selected_file}"
+
+        if file_type == 'mp4':
+            video_file = open(filepath, 'rb')
+            video_bytes = video_file.read()
+            st.video(video_bytes)
+            video_path = filepath
+        else:
+            img = Image.open(filepath).convert('RGB')
+            st.image(img)
+    else:
+        return
+
+
+    
+
+    with st.spinner(f'Analyzing {file_type}...'):
+        if file_type == 'video' or file_type == 'mp4':
+            result = video_pipeline(video_path)
+        else:
+            result = image_pipeline({'images': [img2base64(np.array(img))]})
+
+    if 'found' in result['message']:
+        st.success(result['message'], icon="✅")
+    else:
+        st.error(result['message'], icon="🚨")
+    
+    st.divider()
+    st.write('## Response JSON')
+    st.write(result)
+
+
+def setup():
+
+    if not os.path.isdir("temps"):
+        os.makedirs("temps")
+
+if __name__ == "__main__":
+    image_pipeline = ImagePipeline()
+    video_pipeline = VideoPipeline()
+
+    # with st.sidebar:
+
+    st.title("Face Fake Detection")
+    setup()
+    main()

pipelines.py

@@ -0,0 +1,201 @@
+from PIL import Image
+from ultralytics import YOLO
+
+
+
+from utils import readb64, img2base64
+
+model_int8 = YOLO('weights/best.torchscript', task='detect')
+
+
+def inference_on_image(path):
+    results = model_int8(path)
+
+    img = cv2.imread(path, cv2.COLOR_BGR2RGB)
+    for box in results[0].boxes:
+        cls = box.cls.item()
+        confidence = box.conf.item()
+        label = labels[cls]
+
+        x1, y1, x2, y2 = map(int, list(box.xyxy.numpy()[0]))
+        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 102, 255), 2)
+        img = cv2.rectangle(img, (x1, y1 - 20), (x2, y1), (0, 102, 255), -1)
+        img = cv2.putText(img, "{}: {:.3f}".format(label, confidence), (x1,y1-5),cv2.FONT_HERSHEY_SIMPLEX,0.6,(255,255,255), 1)
+        
+    cv2.imshow('Detected Image', img)
+    cv2.waitKey(0)
+
+    return results
+
+def inference_on_video(path, vid_stride=10):
+    results = model_int8(path, vid_stride=10, stream=True)
+
+    cap = cv2.VideoCapture(path)
+    ret, img = cap.read()
+
+    frame_counter = 0
+    while True:
+        ret, img = cap.read()
+        if ret:
+            if frame_counter % 10 == 0: 
+                result = next(results)
+            for box in result.boxes:
+                cls = box.cls.item()
+                confidence = box.conf.item()
+                label = labels[cls]
+
+                x1, y1, x2, y2 = map(int, list(box.xyxy.numpy()[0]))
+                cv2.rectangle(img, (x1, y1), (x2, y2), (0, 102, 255), 2)
+                img = cv2.rectangle(img, (x1, y1 - 20), (x2, y1), (0, 102, 255), -1)
+                img = cv2.putText(img, "{}: {:.3f}".format(label, confidence), (x1,y1-5),cv2.FONT_HERSHEY_SIMPLEX,0.6,(255,255,255), 1)
+        else:
+            cap.release()
+            break
+        
+        cv2.imshow('Detected Image', img)
+        frame_counter += 1
+
+        k = cv2.waitKey(5) & 0xFF
+        if k == 27:
+            cap.release()
+            cv2.destroyAllWindows()
+            break
+
+    return results
+
+
+class ImagePipeline:
+    def __init__(self, device='cpu', gpu_id=0, weights='weights/best.torchscript'):
+        self.model = YOLO(weights, task='detect')
+
+    def preprocess(self, data):
+        image_base64 = data.pop("images", data)
+
+        if not type(image_base64) == list:
+            image_base64 = [image_base64]
+        elif len(image_base64) > 1:
+            raise Exception("ImagePipeline only accepts 1 image/frame")
+        
+        images = [readb64(image) for image in image_base64]
+        return images
+
+    def inference(self, images):
+        results = self.model(images[0])
+        return results
+
+    def get_response(self, inference_result):
+        response = []
+
+        if 0 in inference_result[0].boxes.cls.numpy():
+            message = "An ambulance is found "
+        else:
+            message = "There is no ambulance"
+
+        for i, result in enumerate(inference_result):
+            for xywhn, cls, conf in zip(
+                result.boxes.xywhn,
+                result.boxes.cls,
+                result.boxes.conf
+            ):
+                xywhn = list(xywhn.numpy())
+                response.append({
+                    'xywhn': {
+                        'x': float(xywhn[0]),
+                        'y': float(xywhn[1]),
+                        'w': float(xywhn[2]),
+                        'h': float(xywhn[3]),
+                    },
+                    'class': cls.item(),
+                    'confidence': conf.item(),
+                })
+
+        return {'results': response,
+                'message': message}
+    
+    def __call__(self, data, config_payload=None):
+        images = self.preprocess(data)
+        inference_result = self.inference(images)
+        response = self.get_response(inference_result)
+        return response
+
+class VideoPipeline:
+    def __init__(self, device='cpu', gpu_id=0, weights='weights/best.torchscript'):
+        self.model = YOLO(weights, task='detect')
+
+    def preprocess(self, data):
+        return data
+
+    def inference(self, video_path, vid_stride=30):
+        results = self.model(video_path, vid_stride=vid_stride)
+        return results
+
+    def get_response(self, inference_result):
+        response = []
+
+        
+        # default message
+        message = "There is no ambulance"
+
+        for i, result in enumerate(inference_result):
+            
+            if 0 in result.boxes.cls.numpy():
+                message = "An ambulance is found"
+
+            for xywhn, cls, conf in zip(
+                result.boxes.xywhn,
+                result.boxes.cls,
+                result.boxes.conf
+            ):
+                xywhn = list(xywhn.numpy())
+                response.append({
+                    'xywhn': {
+                        'x': float(xywhn[0]),
+                        'y': float(xywhn[1]),
+                        'w': float(xywhn[2]),
+                        'h': float(xywhn[3]),
+                    },
+                    'class': cls.item(),
+                    'confidence': conf.item(),
+                })
+
+        return {'results': response,
+                'message': message}
+    
+    def __call__(self, data, config_payload=None):
+        data = self.preprocess(data)
+        inference_result = self.inference(data)
+        response = self.get_response(inference_result)
+        return response
+
+
+if __name__ == '__main__':
+    import cv2
+    import argparse
+ 
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument('--input_type',
+                    default='image',
+                    const='image',
+                    nargs='?',
+                    choices=['image', 'video'],
+                    help='type of input (default: %(default)s)')
+    parser.add_argument("-p", "--path", help="filepath")
+    args = parser.parse_args()
+
+    labels = {
+        0: 'ambulance',
+        1: 'truck'
+    }
+
+    if args.input_type=='image':
+        results = inference_on_image(args.path)
+    elif args.input_type == 'video':
+        results = inference_on_video(args.path)
+    
+    
+    print(results)
+
+    
+    # Examples
+    # python pipelines.py --input_type image --path sample_files/ambulance-2.jpeg
+    # python pipelines.py --input_type video --path sample_files/ambulance.mp4

requirements.txt

@@ -0,0 +1,5 @@
+torch==2.2.1
+ultralytics==8.0.186
+streamlit
+pillow==10.2.0
+opencv-python-headless

sample_files/ambulance.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2f69b30105bad124cc434152707ea3d3ff238f4a9fcb8f2b3774cb2195d5a52c
+size 22860129

sample_files/trucks.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:072f5936821029eab4a1a4866a67df872bfa5b4fa1a51ca7793a4e787435901c
+size 7029598

utils.py

@@ -0,0 +1,101 @@
+import torch
+import numpy as np
+import cv2
+import tempfile, base64
+
+
+def readb64(uri):
+    encoded_data = uri.split(',')[-1]
+    nparr = np.frombuffer(base64.b64decode(encoded_data), np.uint8)
+    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    return img
+
+def img2base64(img, extension="jpg"):
+    _, img_encoded = cv2.imencode(f".{extension}", img)
+    img_base64 = base64.b64encode(img_encoded)
+    img_base64 = img_base64.decode('utf-8')
+    return img_base64
+
+def binary2video(video_binary):
+    # byte_arr = BytesIO()
+    # byte_arr.write(video_binary)
+
+    temp_ = tempfile.NamedTemporaryFile(suffix='.mp4')
+    # decoded_string = base64.b64decode(video_binary)
+
+    temp_.write(video_binary)
+    video_capture = cv2.VideoCapture(temp_.name)
+    ret, frame = video_capture.read()
+    return video_capture
+
+def extract_frames(data_path, interval=30, max_frames=50):
+    """Method to extract frames"""
+    cap = cv2.VideoCapture(data_path)
+    frame_num = 0
+    frames = list()
+
+    while cap.isOpened():
+        success, image = cap.read()
+        if not success:
+            break
+        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        # image = torch.tensor(image) - torch.tensor([104, 117, 123])
+        if frame_num % interval == 0:
+            frames.append(image)
+        frame_num += 1
+        if len(frames) > max_frames:
+            break
+    cap.release()
+    # if len(frames) > max_frames:
+    #     samples = np.random.choice(
+    #         np.arange(0, len(frames)), size=max_frames, replace=False)
+    #     return [frames[_] for _ in samples]
+    return frames
+
+"""FilePicker for streamlit. 
+Still doesn't seem to be a good solution for a way to select files to process from the server Streamlit is running on.
+Here's a pretty functional solution. 
+Usage:
+```
+import streamlit as st
+from filepicker import st_file_selector
+tif_file = st_file_selector(st, key = 'tif', label = 'Choose tif file')
+```
+"""
+
+import os
+import streamlit as st
+
+def update_dir(key):
+    choice = st.session_state[key]
+    if os.path.isdir(os.path.join(st.session_state[key+'curr_dir'], choice)):
+        st.session_state[key+'curr_dir'] = os.path.normpath(os.path.join(st.session_state[key+'curr_dir'], choice))
+        files = sorted(os.listdir(st.session_state[key+'curr_dir']))
+        if "images" in files:
+          files.remove("images")
+        st.session_state[key+'files'] = files
+
+def st_file_selector(st_placeholder, path='.', label='Select a file/folder', key = 'selected'):
+    if key+'curr_dir' not in st.session_state:
+        base_path = '.' if path is None or path == '' else path
+        base_path = base_path if os.path.isdir(base_path) else os.path.dirname(base_path)
+        base_path = '.' if base_path is None or base_path == '' else base_path
+
+        files = sorted(os.listdir(base_path))
+        files.insert(0, 'Choose a file...')
+        if "images" in files:
+          files.remove("images")
+        st.session_state[key+'files'] = files
+        st.session_state[key+'curr_dir'] = base_path
+    else:
+        base_path = st.session_state[key+'curr_dir']
+
+    selected_file = st_placeholder.selectbox(label=label, 
+                                        options=st.session_state[key+'files'], 
+                                        key=key, 
+                                        on_change = lambda: update_dir(key))
+    
+    if selected_file == "Choose a file...":
+        return None
+
+    return selected_file