LightClearNet
PyTorch
Model Details
Image dehazing is a crucial low-level vision task that recovers clear images from hazy ones. In recent years, immense progress has been made in recovering the original image as accurately as possible, without any consideration for efficiency and performance on low-end hardware, making it a mere relic for image beautification. Image dehazing could be very critical for enhanced visibility during indoor emergencies like fire hazards, utilization in surveillance and security cameras and in autonomous vehicles, and hence should be quick and easy to compute in order to be accessible. LightClearNet is a lightweight yet just as effective model for single image dehazing based around batch normalization and skip connections.
Data Overview
The model has been trained on RESIDE-6K.
Improvements
- Rewrote the dehazing logic without using PIL for better performance.
- Integrated PyTube to use any YouTube video without even downloading.
- Quick display and removal of frames feels like running a video natively (work-around as
cv2.imshow()was not working).
To-do
- Listen to keyboard input to control video's flow.
- Implement frame skipping during video playback to maintain a stable video FPS when the number of frames getting dehazed per second is lower.
!pip list | grep opencv
WARNING: No metadata found in /opt/conda/lib/python3.10/site-packages
opencv-contrib-python 4.9.0.80
opencv-python 4.9.0.80
opencv-python-headless 4.9.0.80
!pip install pytube
Collecting pytube
Downloading pytube-15.0.0-py3-none-any.whl.metadata (5.0 kB)
Downloading pytube-15.0.0-py3-none-any.whl (57 kB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 57.6/57.6 kB 2.3 MB/s eta 0:00:00
WARNING: Error parsing requirements for aiohttp: [Errno 2] No such file or directory: '/opt/conda/lib/python3.10/site-packages/aiohttp-3.9.1.dist-info/METADATA'
Installing collected packages: pytube
Successfully installed pytube-15.0.0
import torch
from torchvision import transforms
import cv2
import numpy as np
from pytube import YouTube
from IPython.display import display, clear_output, Image
from timeit import default_timer as timer
import time
from tqdm.auto import tqdm
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device
def cv2_frame_to_img(frame):
return Image(cv2.imencode('.jpg', frame)[1])
def resize_frame(frame, height):
h, w, c = frame.shape
scale = height / float(h)
width = int(w * scale)
return cv2.resize(frame, (width, height))
def frame_dehazer(frame, resize=False):
if resize:
frame = resize_frame(frame, height=400)
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB).astype(np.float32)
frame_rgb /= 255.0
tensor = transforms.ToTensor()(frame_rgb).unsqueeze(0).to(device)
model.eval()
with torch.inference_mode():
pred = model(tensor).squeeze(0)
pred = pred.permute(1, 2, 0)
dehazed_frame = pred.cpu().numpy()
dehazed_frame = (dehazed_frame * 255.0).astype(np.uint8)
dehazed_frame = cv2.cvtColor(dehazed_frame, cv2.COLOR_BGR2RGB)
new_frame = cv2.hconcat([frame, dehazed_frame])
return new_frame
yt_urls = {
'dark_forest': 'https://www.youtube.com/watch?v=LVanPoKh7Ew',
'autumn_forest': 'https://www.youtube.com/watch?v=USa5se_T7TE',
'bright_forest': 'https://www.youtube.com/watch?v=DhsKd0Ln0M8',
'short_clip': 'https://www.youtube.com/watch?v=SgFkqkB6nvM',
'light_fog': 'https://www.youtube.com/watch?v=-pIkWiNJikw',
'almost_still_vid_1': 'https://www.youtube.com/watch?v=8plwv25NYRo',
'almost_still_vid_2': 'https://www.youtube.com/watch?v=yCykHHqk_z0',
}
yt_urls
Out[7]:
{'dark_forest': 'https://www.youtube.com/watch?v=LVanPoKh7Ew',
'autumn_forest': 'https://www.youtube.com/watch?v=USa5se_T7TE',
'bright_forest': 'https://www.youtube.com/watch?v=DhsKd0Ln0M8',
'short_clip': 'https://www.youtube.com/watch?v=SgFkqkB6nvM',
'light_fog': 'https://www.youtube.com/watch?v=-pIkWiNJikw',
'almost_still_vid_1': 'https://www.youtube.com/watch?v=8plwv25NYRo',
'almost_still_vid_2': 'https://www.youtube.com/watch?v=yCykHHqk_z0'}
In [8]:
url = yt_urls['almost_still_vid_2']
yt = YouTube(url)
stream = yt.streams.get_highest_resolution()
# stream = yt.streams.filter(res='480p').first()
# stream = yt.streams.first()
stream
Out[8]:
<Stream: itag="22" mime_type="video/mp4" res="720p" fps="25fps" vcodec="avc1.64001F" acodec="mp4a.40.2" progressive="True" type="video">
In [9]:
stream.title, stream.default_filename, stream.resolution, stream.fps
# stream.download()
Out[9]:
('Winter Wind Sounds 1 Hour / High Winds Blowing Snow',
'Winter Wind Sounds 1 Hour High Winds Blowing Snow.mp4',
'720p',
25)
In [10]:
linkcode
stream_fps = stream.fps
stream_url = stream.url
stream_fps, stream_url
Model Files
cap = cv2.VideoCapture(stream_url)
RESIZE = True
start_after_secs = 80
init_pbar = start_after_secs != 0
total_frames_to_skip = start_after_secs * stream_fps
early_stop, stop_after_secs = True, 30
fps = stream_fps
frames = 0
while True:
start_time = timer()
ret, frame = cap.read()
if not ret:
print('No frames returned.')
break
frames += 1
if init_pbar:
pbar = tqdm(total=start_after_secs, desc='Skipping Start')
init_pbar = False
if frames < total_frames_to_skip:
if frames % stream_fps == stream_fps-1:
pbar.update(1)
continue
dehazed_frame = frame_dehazer(frame, resize=RESIZE)
cv2.putText(
dehazed_frame,
text = f'FPS: {fps:.2f} | Video Time: {round(frames / stream_fps, 2)}s',
org = (10, 20),
fontFace = cv2.FONT_HERSHEY_SIMPLEX,
fontScale = 0.5,
color = (255, 255, 255), # white
thickness = 2
)
display(cv2_frame_to_img(dehazed_frame))
clear_output(wait=True)
end_time = timer()
processing_time = end_time - start_time
sleep_time = (stream_fps ** -1) - processing_time
time.sleep(max(sleep_time, 0))
total_time = timer() - start_time
fps = total_time ** -1
if early_stop and frames > (stop_after_secs * stream_fps) + total_frames_to_skip:
break
cap.release()
0 comments
No comments yet.