From 30dd19772c074adebecdb657c9ebeb218817efcd Mon Sep 17 00:00:00 2001
From: Oldrich Krupauer <oldrich.krupauer@pm.me>
Date: Sat, 19 Oct 2024 11:28:27 +0200
Subject: [PATCH] Changed reflection appearance and added saving images from
 epoch

---
 python_trainer/trainer.py                     | 12 +++++++++++
 .../Assets/Materials/Deepracer - terrain.mat  |  8 ++++++-
 .../Assets/Materials/Road/Amazon road.mat     |  3 +++
 .../Assets/Materials/Road/Reflection.shader   | 21 ++++++++++++++-----
 4 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/python_trainer/trainer.py b/python_trainer/trainer.py
index 5c340b43..eb381ff7 100644
--- a/python_trainer/trainer.py
+++ b/python_trainer/trainer.py
@@ -1,5 +1,6 @@
 import copy
 
+import cv2
 import numpy as np
 import torch
 import torch.onnx
@@ -69,6 +70,17 @@ def train(self, env: UnityEnvironment, temperature: float) -> float:
         if self.writer is not None:
             self.writer.add_image("Sample image", sample_image)
 
+            # Turn to True, if you want to see all of the images from the epoch
+            if False:
+                index = 0
+                for buffer in self.memory.buffer:
+                    for sample in buffer.observations:
+                        gray_image = cv2.cvtColor(
+                            np.resize(sample[0], (100, 160, 1)) * 255.0, cv2.COLOR_GRAY2RGB
+                        )
+                        cv2.imwrite(f"/tmp/{self.curr_epoch}-{index}.jpg", gray_image)
+                        index += 1
+
             # Add text
             steer = ""
 
diff --git a/unity_env/Assets/Materials/Deepracer - terrain.mat b/unity_env/Assets/Materials/Deepracer - terrain.mat
index 820a9aff..395d4104 100644
--- a/unity_env/Assets/Materials/Deepracer - terrain.mat	
+++ b/unity_env/Assets/Materials/Deepracer - terrain.mat	
@@ -8,7 +8,7 @@ Material:
   m_PrefabInstance: {fileID: 0}
   m_PrefabAsset: {fileID: 0}
   m_Name: Deepracer - terrain
-  m_Shader: {fileID: 46, guid: 0000000000000000f000000000000000, type: 0}
+  m_Shader: {fileID: 4800000, guid: 7bd43790af56d1df988076f434fb05d3, type: 3}
   m_Parent: {fileID: 0}
   m_ModifiedSerializedProperties: 0
   m_ValidKeywords: []
@@ -70,14 +70,20 @@ Material:
     - _GlossyReflections: 1
     - _Metallic: 0
     - _Mode: 0
+    - _NoiseScale: 0.07
     - _OcclusionStrength: 1
     - _Parallax: 0.02
+    - _ReflectionStrength: 4
+    - _SetTexture: 0
     - _SmoothnessTextureChannel: 0
     - _SpecularHighlights: 1
+    - _Speed: 0.6
     - _SrcBlend: 1
     - _UVSec: 0
     - _ZWrite: 1
     m_Colors:
     - _Color: {r: 0, g: 0.81960785, b: 0.5294118, a: 1}
     - _EmissionColor: {r: 0, g: 0, b: 0, a: 1}
+    - _MainColor: {r: 0, g: 0.81960785, b: 0.5294118, a: 1}
+    - _ReflectionColor: {r: 1, g: 0.9411765, b: 0.6117647, a: 1}
   m_BuildTextureStacks: []
diff --git a/unity_env/Assets/Materials/Road/Amazon road.mat b/unity_env/Assets/Materials/Road/Amazon road.mat
index f9a92a04..d5784bd5 100644
--- a/unity_env/Assets/Materials/Road/Amazon road.mat	
+++ b/unity_env/Assets/Materials/Road/Amazon road.mat	
@@ -84,6 +84,7 @@ Material:
     - _OcclusionStrength: 1
     - _Parallax: 0.02
     - _ReflectionStrength: 4
+    - _SetTexture: 1
     - _Smoothness: 0
     - _SmoothnessTextureChannel: 1
     - _SpecularHighlights: 1
@@ -94,5 +95,7 @@ Material:
     m_Colors:
     - _Color: {r: 1, g: 1, b: 1, a: 1}
     - _EmissionColor: {r: 0, g: 0, b: 0, a: 1}
+    - _MainColor: {r: 1, g: 1, b: 0, a: 1}
+    - _ReflectionColor: {r: 1, g: 0.9415832, b: 0.6133826, a: 1}
     - _SpecColor: {r: 0.19999996, g: 0.19999996, b: 0.19999996, a: 1}
   m_BuildTextureStacks: []
diff --git a/unity_env/Assets/Materials/Road/Reflection.shader b/unity_env/Assets/Materials/Road/Reflection.shader
index e43e5efd..b78f2396 100644
--- a/unity_env/Assets/Materials/Road/Reflection.shader
+++ b/unity_env/Assets/Materials/Road/Reflection.shader
@@ -3,9 +3,13 @@ Shader "Custom/Reflection"
     Properties
     {
         _MainTex ("Albedo (RGB)", 2D) = "white" {}
+        [Maincolor] _MainColor ("Main color", Color) = (1., 1., 0., 1.)
+        _SetTexture ("Texture set", Float) = 0.
         _ReflectionStrength ("Reflection strength", Float) = 1.0
-        _NoiseScale ("Noise scale", Float) = 1.0
+        _NoiseScaleX ("Noise scale X", Float) = 1.0
+        _NoiseScaleY ("Noise scale Y", Float) = 1.0
         _Speed ("Noise speed", Range(0, 1)) = 0.5
+        _ReflectionColor ("Reflection color", Color) = (1., 1., 1., 1.)
     }
     SubShader
     {
@@ -20,9 +24,13 @@ Shader "Custom/Reflection"
         #pragma target 3.0
 
         sampler2D _MainTex;
+        float4 _MainColor;
         float _ReflectionStrength;
-        float _NoiseScale;
+        float _NoiseScaleX;
+        float _NoiseScaleY;
         float _Speed;
+        float4 _ReflectionColor;
+        float _SetTexture;
 
         struct Input
         {
@@ -166,12 +174,15 @@ Shader "Custom/Reflection"
         void surf (Input IN, inout SurfaceOutputStandard o)
         {
             // float2 noiseInput = IN.uv_MainTex * _NoiseScale;
-            float2 noiseInput = IN.worldPos.xz * _NoiseScale;
+            float2 noiseInput = IN.worldPos.xz;
+            noiseInput.x *= _NoiseScaleX;
+            noiseInput.y *= _NoiseScaleY;
             float noiseValue = clamp(ClassicNoise(float3(noiseInput.x, noiseInput.y, _Time.y * _Speed)), 0.0, 1.0) * _ReflectionStrength;
             // float noiseValue = ClassicNoise(noiseInput);
 
-            fixed4 albedo = tex2D(_MainTex, IN.uv_MainTex);
-            o.Albedo = albedo + noiseValue;
+            fixed4 albedo = tex2D(_MainTex, IN.uv_MainTex) * _SetTexture + _MainColor * (1. - _SetTexture);
+
+            o.Albedo = albedo + noiseValue * _ReflectionColor.rgb;
         }
         ENDCG
     }